package com.xshang.tenant;

import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;

/**
 * 多租户数据源
 */
@Slf4j
public class TenantDataSource extends AbstractRoutingDataSource {
    /**
     * 租户数据源集合：key为租户id或唯一编码，value为具体租户的数据源
     */
    private final Map<Object, Object> targetDataSources;

    /**
     * 多租户数据源构造器
     *
     * @param defaultDataSource 默认数据源
     * @param targetDataSources 租户数据源集合
     */
    public TenantDataSource(DataSource defaultDataSource, Map<Object, Object> targetDataSources){
        this.targetDataSources = targetDataSources;
        super.setDefaultTargetDataSource(defaultDataSource);
        super.setTargetDataSources(this.targetDataSources);
        // afterPropertiesSet负责解析成可用的目标数据源
        super.afterPropertiesSet();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        // 通过线程上下文设置方法获取当前租户的数据源名称
        return TenantDataSourceContext.get();
    }

    /**
     * 添加数据源到租户数据源集合
     *
     * @param dataSourceName 数据源名称
     * @param dataSource 数据源
     */
    public void add(String dataSourceName, DataSource dataSource) {
        if (targetDataSources.containsKey(dataSourceName)) {
            log.warn("新增数据源: {} 数据源已存在，请检查。", dataSourceName);
            return;
        }
        targetDataSources.put(dataSourceName, dataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 更新租户数据源集合中的数据源
     *
     * @param dataSourceName 数据源名称
     * @param dataSource 数据源
     */
    public void update(String dataSourceName, DataSource dataSource) {
        if (!targetDataSources.containsKey(dataSourceName)) {
            log.warn("更新数据源: {} 找不到数据源，已忽略", dataSourceName);
            return;
        }
        targetDataSources.replace(dataSourceName, dataSource);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 从租户数据源集合中移除指定的数据源
     *
     * @param dataSourceName 数据源名称
     */
    public void remove(String dataSourceName) {
        // 缺省租户数据源不允许删除
        if (TenantDataSourceContext.DEFAULT_TENANT_DATA_SOURCE.equals(dataSourceName)) {
            log.warn("删除数据源：defaultTenantDataSource 不允许删除");
            return;
        }
        if (!targetDataSources.containsKey(dataSourceName)) {
            log.warn("删除数据源: {} 数据源不存在, 已忽略.", dataSourceName);
            return;
        }
        targetDataSources.remove(dataSourceName);
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
}
